#!/usr/bin/env bash
set -uo pipefail
#==============================================================#
# File      :   pg-restore
# Desc      :   script to restore pg cluster with pgbackrest
# Time      :   {{ '%Y-%m-%d %H:%M' | strftime }}
# Host      :   {{ pg_instance }} @ {{ inventory_hostname }}:{{ patroni_port }}
# Path      :   /pg/bin/pg-restore
# Deps      :   /usr/bin/pgbackrest, /pg/conf/pitr.conf
# License   :   AGPLv3 @ https://doc.pgsty.com/about/license
# Copyright :   2018-2025  Ruohang Feng / Vonng (rh@vonng.com)
#==============================================================#
# https://pgbackrest.org/command.html#command-restore
# https://pgbackrest.org/configuration.html#section-restore
# https://www.postgresql.org/docs/current/runtime-config-wal.html#RUNTIME-CONFIG-WAL-RECOVERY-TARGET

#--------------------------------------------------------------#
# Utils
#--------------------------------------------------------------#
__CN='\033[0m';__CB='\033[0;30m';__CR='\033[0;31m';__CG='\033[0;32m';
__CY='\033[0;33m';__CB='\033[0;34m';__CM='\033[0;35m';__CC='\033[0;36m';__CW='\033[0;37m';
function log_info() {  printf "${__CG}$*${__CN}\n";   }
function log_warn() {  printf "${__CY}$*${__CN}\n";   }
function log_error() { printf "${__CR}$*${__CN}\n";   }
function log_debug() { printf "${__CB}$*${__CN}\n"; }
function log_input() { printf "${__CM}$*\n=> ${__CN}"; }
function log_hint()  { printf "${__CB}$*${__CN}\n"; }
function log_line()  { printf "${__CM}[$*] ===========================================${__CN}\n"; }
function log_error() { printf "[${__CR}FAIL${__CN}] ${__CR}$*${__CN}\n"; exit 1;  }

{% macro abort(error) %}{{ None['[ERROR] ' ~ error][0] }}{% endmacro %}

{% set p = pg_pitr                   | default({})           %}
{% set pitr_repo       = p.repo      | default(pgbackrest_repo[pgbackrest_method]) %}
{% set pitr_cluster    = p.cluster   | default(pg_cluster) | string %}
{% set pitr_type       = p.type      | default('default')  | string %}
{% set pitr_action     = p.action    | default('pause')    | string %}
{% set pitr_set        = p.set       | default('latest')   | string %}
{% set pitr_timeline   = p.timeline  | default('latest')   | string %}
{% set pitr_data       = p.data      | default(pg_data)    | string %}

{% set cmd_args = [] %}
{% set cmd_args = cmd_args + [ '--config=/pg/conf/pitr.conf' ] %}
{% set cmd_args = cmd_args + [ '--stanza=' + pitr_cluster    ] %}

{% if pitr_type not in ['default','immediate','time','lsn','xid','name'] %}
{{ abort("invalid pg_pitr.type: " + pitr_type) }}
{% endif %}

{% if p.time is defined %}
# pitr type = time (recover to the time point {{ p.time }})
{% if p.lsn is defined or p.xid is defined or p.name is defined or p.type|default('time') != 'time' %}
{{ abort("invalid pitr target provided!") }}
{% endif %}
# validate time format, fill to YYYY-MM-DD HH24:MI:SS.US+00 (if timezone not given, fill with current timezone in +xx/-xx format)
# if timezone not given, fill with current timezone in +xx/-xx format
{% set pitr_type = 'time' %}
{% set pitr_target = p.time %}
{% set cmd_args = cmd_args + ['--type=time'] %}
{% set cmd_args = cmd_args + ['--target="' + p.time + '"'] %}
{% endif %}

{% if p.name is defined %}
# pitr type = name (recover to the restore point {{ p.name }} created by pg_create_restore_point)
{% if p.lsn is defined or p.xid is defined or p.time is defined or p.type|default('name') != 'name' %}
{{ abort("invalid pitr target provided!") }}
{% endif %}
{% set pitr_type = 'name' %}
{% set cmd_args = cmd_args + ['--type=name'] %}
{% set cmd_args = cmd_args + ['--target="' + p.name + '"'] %}
{% endif %}

{% if p.xid is defined %}
# pitr type = xid (recover to the transaction id {{ p.xid }})
{% if p.lsn is defined or p.name is defined or p.time is defined or p.type|default('xid') != 'xid' %}
{{ abort("invalid pitr target provided!") }}
{% endif %}
{% if p.xid|int is not number or p.xid|int <= 0 or p.xid|int >= 4294967296 %}
{{ abort("invalid XID: " + p.xid + ". XID must be a positive 32-bit integer (0 < xid < 4294967296)") }}
{% endif %}
{% set pitr_type = 'xid' %}
{% set cmd_args = cmd_args + ['--type=xid'] %}
{% set cmd_args = cmd_args + ['--target=' + p.xid|string ] %}
{% endif %}

{% if p.lsn is defined %}
# pitr type = lsn (recover to the LSN point {{ p.lsn }})
{% if p.xid is defined or p.name is defined or p.time is defined or p.type|default('lsn') != 'lsn' %}
{{ abort("invalid pitr target provided!") }}
{% endif %}
{% if not p.lsn|upper is match("[0-9A-F]{1,8}/[0-9A-F]{1,8}") %}
{{ abort("invalid LSN format: " + p.lsn + ". Expected format: [0-9A-F]{1,8}/[0-9A-F]{1,8}") }}
{% endif %}
{% set pitr_type = 'lsn' %}
{% set cmd_args = cmd_args + ['--type=lsn'] %}
{% set cmd_args = cmd_args + ['--target=' + p.lsn|upper] %}
{% endif %}

{% if p.type is defined and p.type == 'immediate' %}
# pitr type = immediate ( recover only until the database becomes consistent since {{ pitr_set }} backup)
{% if p.xid is defined or p.name is defined or p.time is defined or p.lsn is defined  %}
{{ abort("invalid pitr target provided!") }}
{% endif %}
{% set pitr_type = 'immediate' %}
{% set cmd_args = cmd_args + ['--type=immediate'] %}
{% endif %}

{% if p.time is not defined and p.xid is not defined and p.name is not defined and p.lsn is not defined and (p.type is not defined or p.type == 'default') %}
# pitr type = default (recover to the end of the archive stream.)
{% set pitr_type = 'default' %}
{% endif %}

{% if p.exclusive is defined and p.exclusive|bool %}
{% set cmd_args = cmd_args + ['--target-exclusive'] %}
{% endif %}

{% if pitr_action not in ['promote','pause','shutdown'] %}
{{ abort("invalid pg_pitr.action: " + pitr_action) }}
{% else %}
{% if pitr_action != 'pause' and pitr_type in ['immediate','time','name','xid','lsn'] %}
# target action = {{ pitr_action }} when {{ pitr_type }} pitr is done
{% set cmd_args = cmd_args + ['--target-action=' + pitr_action] %}
{% endif %}
{% endif %}

{% if pitr_timeline|string != 'latest' %}
# timeline is specified as {{ pitr_timeline }}
{% set cmd_args = cmd_args + ['--target-timeline=' + pitr_timeline|string] %}
{% endif %}

{% if pitr_set != 'latest' %}
# restore since the specified backup set {{ pitr_set }}
{% set cmd_args = cmd_args + ['--set=' + pitr_set] %}
{% endif %}

# run this script as postgres dbsu
DBSU={{ pg_dbsu|default('postgres') }}
if [ "$(whoami)" != "$DBSU" ]; then
    log_error "This script must be run as $DBSU os user"
    exit 1
fi




log_warn "Perform {{ pitr_type }} PITR on {{ pg_cluster }} from stanza {{ pitr_cluster }}"

log_line "1. Stop PostgreSQL"
log_info "   1.1 Pause Patroni (if there are any replicas)"
log_hint "       $ pg pause <cls>  # pause patroni auto failover"
log_info "   1.2 Shutdown Patroni"
log_hint "       $ pt-stop         # sudo systemctl stop patroni"
log_info "   1.3 Shutdown Postgres"
log_hint "       $ pg-stop         # pg_ctl -D {{ pitr_data }} stop -m immediate"
log_hint ""

log_line "2. Perform PITR"
log_info "   2.1 Restore Backup"
log_hint "       $ " /usr/bin/pgbackrest {{ cmd_args|join(' ') }} restore
log_info "   2.2 Start PG to Replay WAL"
log_hint "       $ pg-start        # pg_ctl -D {{ pitr_data }} start"
log_info "   2.3 Validate and Continue"
log_warn "     - If it is not the point you want, goto 2.1"
log_info "     - Otherwise make sure pg-promote then goto 3"
log_hint ""

log_line "3. Restore Primary"
log_info "   3.1 Enable Archive Mode (Restart Required)"
log_hint "       $ psql -c 'ALTER SYSTEM SET archive_mode = on;'"
log_info "   3.1 Restart Postgres to Apply Changes"
log_hint "       $ pg-restart      # pg_ctl -D {{ pitr_data }} restart"
log_info "   3.3 Restart Patroni"
log_hint "       $ pt-restart      # sudo systemctl restart patroni"
log_hint ""


log_warn "WARNING: You are about to perform a Point-In-Time Recovery!"
log_warn "This operation will nuke the existing cluster with new one!"
log_warn "If this is not what you want, Ctrl+C to abort before start!"

# Countdown from 5 to 1
for i in 5 4 3 2 1; do
    log_hint "\rCountdown: $i seconds remaining..."
    sleep 1
done
log_warn "Proceeding with PITR operation."

log_line "Restore Begin!"

log_line "1.1 Pause Patroni"
patronictl -c /pg/bin/patroni.yml pause "{{ pg_cluster }}"

log_line "1.2 Shutdown Patroni"
sudo systemctl stop patroni

log_line "1.3 Shutdown Postgres"
pg_ctl -D "{{ pitr_data }}" stop

log_line "2.1 Perform Restore"
log_hint "$ /usr/bin/pgbackrest" {{ cmd_args|join(' ') }} restore
/usr/bin/pgbackrest {{ cmd_args|join(' ') }} restore

log_line "2.2 Start PG to Replay WAL"
pg_ctl -D "{{ pitr_data }}" start

log_line "Restore Done!"

log_line "What's next?"
log_info "   2.3 Validate Your backup, then make sure promoted "
log_hint "       $ pg-promote"
log_info "   3.1 Enable Archive Mode (Restart Required)"
log_hint "       $ psql -c 'ALTER SYSTEM SET archive_mode = on;'"
log_info "   3.1 Restart Postgres to Apply Changes"
log_hint "       $ pg-restart      # pg_ctl -D {{ pitr_data }} restart"
log_info "   3.3 Restart Patroni"
log_hint "       $ pt-restart      # sudo systemctl restart patroni"
log_info "   4.x Reinit other replicas, if you have"
